"""
可视化法线分布
"""
import numpy as np
import taichi as ti

from pyrender.gmath import *  # pylint: disable=wildcard-import,unused-wildcard-import
from pyrender.gtaichi import * # pylint: disable=wildcard-import,unused-wildcard-import
from pyrender.scenes import * # pylint: disable=wildcard-import,unused-wildcard-import
from pyrender.render_frame import RendererFrame

ti.init()
N = 16000 #支持的最大对象数目

@ti.data_oriented
class NormalRenderer:
    """
    渲染法线分布
    """
    def __init__(self,w=None,h=None,scene=None):
        self.resolution = (800 if w is None else w,450 if h is None else h)
        self.pixels = ti.Vector.field(3, float, self.resolution) # 图像的像素场
        self.mat4 = ti.Matrix.field(4,4,dtype=ti.f32,shape=(1,))
        self.scene = scene
        self.px = ti.Vector.field(3,float,(N,))

    @ti.kernel
    def build_random_normal(self):
        a = .01
        n = ti.Vector([0.,0.,1.])
        ray = ti.Vector([-1.,-1.,-1.])
        for i in range(N):
            diffuse_dir  = random_dir(n)
            specular_dir = ti.math.normalize(ti.math.reflect(ray,n))
            if ti.random() > a:
                self.px[i] = 10. * ti.math.normalize(lerp(specular_dir,diffuse_dir,10*a))
            else:
                self.px[i] = 10. * diffuse_dir

    @ti.kernel
    def build_importance_sample_ggx(self):
        a = 0.1
        n = ti.Vector([0.,0.,1.])
        ray = ti.Vector([-1.,-1.,-1.])
        for i in range(N):
            xi = ti.Vector([ti.random(),ti.random()])
            self.px[i] = 10. * importance_sample_ggx(xi,n,a)

    def build_render_data(self):
        """
        重建渲染数据
        """
        self.update_mat4()
        self.build_random_normal()

    @ti.func
    def fill_screen(self,c:vec3):
        """
        填充屏幕颜色
        """
        for i,j in self.pixels:
            self.pixels[i,j] = c

    def update_mat4(self):
        for i in range(4):
            for j in range(4):
                self.mat4[0][i,j] = self.scene.camera.mat4[i,j]

    def reset_framebuffer(self):
        """
        重置数据
        """
        self.update_mat4()

    def accumulate(self):
        """
        累加渲染
        """
        self.render()

    def fetch_image(self):
        """
        返回渲染images
        """
        return self.pixels

    @ti.func
    def draw_px(self,x,y):
        if x<=1. and x>=0. and y<=1. and y>=0.:
            ax = int(x*(self.resolution[0]-1))
            ay = int(y*(self.resolution[1]-1))

            self.pixels[ax,ay] = ti.Vector([1.,0.,0.])

    @ti.kernel
    def render(self):
        """
        渲染
        """
        f = 2.
        pxx = ti.Vector([0.,0.,0.])
        self.fill_screen(pxx)
        for i in range(N):
            pt = self.px[i]
            px = self.mat4[0] @ ti.Vector([pt.x,pt.y,pt.z,1.])

            pxx.x = f * px.x/(-px.z+f)+.5
            pxx.y = f * px.y/(-px.z+f)+.5
            self.draw_px(pxx.x,pxx.y)

scene = Scenes()
scene.set_camera(Camera(np.array((50.,50.,50.)),np.array((0.,0.,0.))))
render = NormalRenderer(640,640,scene=scene)

RendererFrame(render,1.6).loop()
